{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook is intended to demonstrate how select registration, segmentation, and image mathematical methods of ITKTubeTK can be combined to perform multi-channel brain extraction (aka. skull stripping for patient data containing multiple MRI sequences).\n",
    "\n",
    "There are many other (probably more effective) brain extraction methods available as open-source software such as BET and BET2 in the FSL package (albeit such methods are only for single channel data).   If you need to perform brain extraction for a large collection of scans that do not contain major pathologies, please use one of those packages.   This notebook is meant to show off the capabilities of specific ITKTubeTK methods, not to demonstration how to \"solve\" brain extraction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import itk\n",
    "from itk import TubeTK as ttk\n",
    "\n",
    "from itkwidgets import view\n",
    "\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "InputBaseDir = \"../Data/CTP-MinMax/\"\n",
    "\n",
    "CTPMaxFilename = InputBaseDir + \"max3.mha\"\n",
    "CTPMinFilename = InputBaseDir + \"min3.mha\"\n",
    "CTPBrainFilename = InputBaseDir + \"max3-Brain.mha\"\n",
    "\n",
    "imMax = itk.imread(CTPMaxFilename, itk.F)\n",
    "imMin = itk.imread(CTPMinFilename, itk.F)\n",
    "imBrain = itk.imread(CTPBrainFilename, itk.F)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8a4d0906aa3a4316ae1a7a00d9add363",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Viewer(geometries=[], gradient_opacity=0.22, point_sets=[], rendered_image=<itk.itkImagePython.itkImageF3; pro…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "view(imBrain)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "ImageType = itk.Image[itk.F, 3]\n",
    "\n",
    "imMath = ttk.ImageMath.New(Input=imBrain)\n",
    "imMath.Threshold( 0.00001, 2000, 1, 0)\n",
    "imMath.Erode(10,1,0)\n",
    "imBrainMaskErode = imMath.GetOutput()\n",
    "\n",
    "imMath.SetInput(imMax)\n",
    "imMath.AddImages(imMin,1,-1)\n",
    "imDiff = imMath.GetOutput()\n",
    "imMath.ReplaceValuesOutsideMaskRange(imBrain, 0.0001, 2000, 0)\n",
    "imDiffBrain = imMath.GetOutput()\n",
    "imMath.ReplaceValuesOutsideMaskRange(imBrainMaskErode, 0.5, 1.5, 0)\n",
    "imDiffBrainErode = imMath.GetOutput()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 590.79504\n"
     ]
    }
   ],
   "source": [
    "tmpA = itk.GetArrayViewFromImage(imDiffBrain)\n",
    "tmpAE = itk.GetArrayViewFromImage(imDiffBrainErode)\n",
    "zMax = tmpA.shape[0]\n",
    "clip = 0\n",
    "while((np.amax(tmpA[clip:clip+1,:,:])>1000) | (np.amax(tmpA[clip:clip+1,:,:])==0)):\n",
    "    clip += 1\n",
    "if(clip>0):\n",
    "    tmpA[0:clip,:,:]=0\n",
    "    tmpAE[0:clip,:,:]=0\n",
    "clip = 1\n",
    "while((np.amax(tmpA[zMax-clip:zMax-clip+1,:,:])>1000) | (np.amax(tmpA[zMax-clip:zMax-clip+1,:,:])==0)):\n",
    "    clip += 1\n",
    "print(clip, np.amax(tmpA[zMax-clip:zMax-clip+1,:,:]))\n",
    "clip = clip - 1\n",
    "if(clip>0):\n",
    "    tmpA[zMax-clip:zMax,:,:]=0  #Happens to imDiffBrain since this array is a view of an itk image\n",
    "    tmpAE[zMax-clip:zMax,:,:]=0  #Happens to imDiffBrain since this array is a view of an itk image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "463bd6c9c9114629988b0937d5a3ad81",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Viewer(geometries=[], gradient_opacity=0.22, point_sets=[], rendered_image=<itk.itkImagePython.itkImageF3; pro…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "view(imDiffBrain)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ -11.42652893 -149.88652039  357.57658386]\n",
      " [  -4.23602295 -145.46159363  296.73384094]\n",
      " [ -33.55116272 -160.39572144  298.94630432]\n",
      " [ -52.35710144 -206.85745239  296.73384094]\n",
      " [ -25.80754089 -144.90847778  278.48101807]\n",
      " [ -48.48529053 -193.02955627  305.03057861]\n",
      " [   0.74201965 -167.03311157  264.1000061 ]\n",
      " [ -45.16659546 -175.88296509  317.75224304]\n",
      " [  -5.34225464 -133.29304504  307.24304199]\n",
      " [ -58.44137573 -155.97079468  292.86203003]\n",
      " [  -6.44848633 -155.41767883  264.1000061 ]\n",
      " [ -43.50724792 -140.48355103  327.70832825]\n",
      " [  26.18534851 -133.29304504  285.1184082 ]\n",
      " [ -46.82594299 -157.63014221  276.82167053]\n",
      " [ -18.61703491 -128.86811829  290.64956665]]\n"
     ]
    }
   ],
   "source": [
    "imMath = ttk.ImageMath[ImageType,ImageType].New()\n",
    "imMath.SetInput(imDiffBrainErode)\n",
    "imMath.Blur(1.5)\n",
    "imBlur = imMath.GetOutput()\n",
    "imBlurArray = itk.GetArrayViewFromImage(imBlur)\n",
    "\n",
    "numSeeds = 15\n",
    "seedCoverage = 20\n",
    "seedCoord = np.zeros([numSeeds,3])\n",
    "for i in range(numSeeds):\n",
    "    seedCoord[i] = np.unravel_index(np.argmax(imBlurArray, axis=None), imBlurArray.shape)\n",
    "    indx = [int(seedCoord[i][0]),int(seedCoord[i][1]),int(seedCoord[i][2])]\n",
    "    minX = max(indx[0]-seedCoverage,0)\n",
    "    maxX = max(indx[0]+seedCoverage,imBlurArray.shape[0])\n",
    "    minY = max(indx[1]-seedCoverage,0)\n",
    "    maxY = max(indx[1]+seedCoverage,imBlurArray.shape[1])\n",
    "    minZ = max(indx[2]-seedCoverage,0)\n",
    "    maxZ = max(indx[2]+seedCoverage,imBlurArray.shape[2])\n",
    "    imBlurArray[minX:maxX,minY:maxY,minZ:maxZ]=0\n",
    "    indx.reverse()\n",
    "    seedCoord[:][i] = imDiffBrain.TransformIndexToPhysicalPoint(indx)\n",
    "print(seedCoord)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "**** Processing seed 0 : [ -11.42652893 -149.88652039  357.57658386]\n",
      "**** Processing seed 1 : [  -4.23602295 -145.46159363  296.73384094]\n",
      "**** Processing seed 2 : [ -33.55116272 -160.39572144  298.94630432]\n",
      "**** Processing seed 3 : [ -52.35710144 -206.85745239  296.73384094]\n",
      "**** Processing seed 4 : [ -25.80754089 -144.90847778  278.48101807]\n",
      "**** Processing seed 5 : [ -48.48529053 -193.02955627  305.03057861]\n",
      "**** Processing seed 6 : [   0.74201965 -167.03311157  264.1000061 ]\n",
      "**** Processing seed 7 : [ -45.16659546 -175.88296509  317.75224304]\n",
      "**** Processing seed 8 : [  -5.34225464 -133.29304504  307.24304199]\n",
      "**** Processing seed 9 : [ -58.44137573 -155.97079468  292.86203003]\n",
      "**** Processing seed 10 : [  -6.44848633 -155.41767883  264.1000061 ]\n",
      "**** Processing seed 11 : [ -43.50724792 -140.48355103  327.70832825]\n",
      "**** Processing seed 12 : [  26.18534851 -133.29304504  285.1184082 ]\n",
      "**** Processing seed 13 : [ -46.82594299 -157.63014221  276.82167053]\n",
      "**** Processing seed 14 : [ -18.61703491 -128.86811829  290.64956665]\n"
     ]
    }
   ],
   "source": [
    "# Manually extract a few vessels to form an image-specific training set\n",
    "vSeg = ttk.SegmentTubes.New(Input=imDiffBrain)\n",
    "vSeg.SetVerbose(True)\n",
    "vSeg.SetMinRoundness(0.4)\n",
    "vSeg.SetMinCurvature(0.002)\n",
    "vSeg.SetRadiusInObjectSpace( 1 )\n",
    "for i in range(numSeeds):\n",
    "    print(\"**** Processing seed \" + str(i) + \" : \" + str(seedCoord[i]))\n",
    "    vSeg.ExtractTubeInObjectSpace( seedCoord[i], i )\n",
    "    \n",
    "tubeMaskImage = vSeg.GetTubeMaskImage()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7a45792b565448dfa5e16439bf2b7091",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Viewer(geometries=[], gradient_opacity=0.22, point_sets=[], rendered_image=<itk.itkImagePython.itkImageF3; pro…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "imMath.SetInput(tubeMaskImage)\n",
    "imMath.AddImages(imDiffBrain, 200, 1)\n",
    "blendIm = imMath.GetOutput()\n",
    "view(blendIm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "LabelMapType = itk.Image[itk.UC,3]\n",
    "\n",
    "trMask = ttk.ComputeTrainingMask[ImageType,LabelMapType].New()\n",
    "trMask.SetInput( tubeMaskImage )\n",
    "trMask.SetGap( 4 )\n",
    "trMask.SetObjectWidth( 1 )\n",
    "trMask.SetNotObjectWidth( 1 )\n",
    "trMask.Update()\n",
    "fgMask = trMask.GetOutput()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "406d101cb39e42c6b8e7d3dacfd9d593",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Viewer(geometries=[], gradient_opacity=0.22, point_sets=[], rendered_image=<itk.itkImagePython.itkImageUC3; pr…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "view(fgMask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "enhancer = ttk.EnhanceTubesUsingDiscriminantAnalysis[ImageType,LabelMapType].New()\n",
    "enhancer.AddInput( imDiff )\n",
    "enhancer.SetLabelMap( fgMask )\n",
    "enhancer.SetRidgeId( 255 )\n",
    "enhancer.SetBackgroundId( 128 )\n",
    "enhancer.SetUnknownId( 0 )\n",
    "enhancer.SetTrainClassifier(True)\n",
    "enhancer.SetUseIntensityOnly(True)\n",
    "enhancer.SetScales([0.43,1.29,3.01])\n",
    "enhancer.Update()\n",
    "enhancer.ClassifyImages()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "45e5ababc0544b768eb05f835431d9c2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Viewer(geometries=[], gradient_opacity=0.22, point_sets=[], rendered_image=<itk.itkImagePython.itkImageF3; pro…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "im1vess = itk.SubtractImageFilter( Input1=enhancer.GetClassProbabilityImage(0), Input2=enhancer.GetClassProbabilityImage(1))\n",
    "\n",
    "imMath.SetInput(imDiffBrain)\n",
    "imMath.Threshold(0.0001,2000,1,0)\n",
    "imMath.Erode(2,1,0)\n",
    "imBrainE = imMath.GetOutput()\n",
    "\n",
    "imMath.SetInput(im1vess)\n",
    "imMath.ReplaceValuesOutsideMaskRange(imBrainE, 1, 1, -0.001)\n",
    "im1vessBrain = imMath.GetOutput()\n",
    "#view(enhancer.GetClassProbabilityImage(0))\n",
    "view(im1vessBrain)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "itk.imwrite( im1vess, InputBaseDir + \"diff3-VesselEnhanced.mha\", compression=True)\n",
    "\n",
    "itk.imwrite( im1vessBrain, InputBaseDir + \"diff3-Brain-VesselEnhanced.mha\", compression=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
